import io
from datetime import datetime
from django.utils import timezone
from django.db.models import Sum
import openpyxl
from openpyxl.styles import Font, Alignment
from openpyxl.utils import get_column_letter

class QuerySetToXlsxExporter:
    """
    通用QuerySet导出为XLSX文件类
    """
    
    def __init__(self, queryset, group_by_fields, sum_fields, filename=None):
        """
        初始化导出器
        
        Args:
            queryset: Django QuerySet (model.objects.filter(enabled=True).values()的结果)
            group_by_fields: list, 用于分组的字段列表
            sum_fields: list, 需要求和的字段列表
            filename: str, 导出文件名
        """
        self.queryset = queryset
        self.group_by_fields = group_by_fields
        self.sum_fields = sum_fields
        self.filename = filename or f"export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.xlsx"
        self.workbook = None
        self.ws = None
    
    def _convert_value(self, value):
        """
        数据值转换
        
        Args:
            value: 原始值
            
        Returns:
            转换后的值
        """
        if value is None:
            return ""
        elif isinstance(value, bool):
            return "是" if value else "否"
        elif isinstance(value, datetime):
            # 处理时区时间，转换为本地时间字符串
            if timezone.is_aware(value):
                value = timezone.localtime(value)
            return value.strftime("%Y-%m-%d %H:%M:%S")
        elif hasattr(value, 'strftime'):  # 其他日期类型
            return value.strftime("%Y-%m-%d %H:%M:%S")
        else:
            return value
    
    def _get_aggregated_data(self):
        """
        执行分组聚合查询
        
        Returns:
            聚合后的QuerySet
        """
        # 构建annotate参数
        annotate_params = {
            f"{field}__sum": Sum(field) for field in self.sum_fields
        }
        
        # 执行分组聚合
        aggregated_data = self.queryset.values(*self.group_by_fields).annotate(**annotate_params)
        
        return aggregated_data
    
    def _setup_worksheet(self):
        """设置工作表格式"""
        # 设置标题行样式
        header_font = Font(bold=True, size=12)
        header_alignment = Alignment(horizontal='center', vertical='center')
        
        for col_num, column_title in enumerate(self.column_titles, 1):
            cell = self.ws.cell(row=1, column=col_num, value=column_title)
            cell.font = header_font
            cell.alignment = header_alignment
    
    def _auto_adjust_column_width(self):
        """自动调整列宽"""
        for col_num, column_title in enumerate(self.column_titles, 1):
            max_length = len(str(column_title))
            column_letter = get_column_letter(col_num)
            
            for row_num in range(2, self.ws.max_row + 1):
                cell_value = self.ws.cell(row=row_num, column=col_num).value
                if cell_value:
                    max_length = max(max_length, len(str(cell_value)))
            
            # 设置列宽，留一些余量
            adjusted_width = min(max_length + 2, 50)
            self.ws.column_dimensions[column_letter].width = adjusted_width
    
    def _generate_column_titles(self, aggregated_data):
        """
        生成列标题
        
        Args:
            aggregated_data: 聚合数据
            
        Returns:
            list: 列标题列表
        """
        column_titles = []
        
        # 添加分组字段标题
        for field in self.group_by_fields:
            # 将字段名转换为更友好的显示名称
            display_name = field.replace('_', ' ').title()
            column_titles.append(display_name)
        
        # 添加求和字段标题
        for field in self.sum_fields:
            display_name = f"{field.replace('_', ' ').title()} (合计)"
            column_titles.append(display_name)
        
        # 添加其他非聚合字段
        if aggregated_data:
            sample_item = aggregated_data[0]
            for field in sample_item.keys():
                if (field not in self.group_by_fields and 
                    not field.endswith('__sum') and
                    field not in column_titles):
                    display_name = field.replace('_', ' ').title()
                    column_titles.append(display_name)
        
        return column_titles
    
    def _get_field_mapping(self, aggregated_data):
        """
        生成字段映射关系
        
        Args:
            aggregated_data: 聚合数据
            
        Returns:
            dict: 字段映射关系
        """
        field_mapping = {}
        current_index = 0
        
        # 分组字段映射
        for field in self.group_by_fields:
            field_mapping[field] = current_index
            current_index += 1
        
        # 求和字段映射
        for field in self.sum_fields:
            field_mapping[f"{field}__sum"] = current_index
            current_index += 1
        
        # 其他字段映射
        if aggregated_data:
            sample_item = aggregated_data[0]
            for field in sample_item.keys():
                if (field not in field_mapping and 
                    not field.endswith('__sum')):
                    field_mapping[field] = current_index
                    current_index += 1
        
        return field_mapping
    
    def export_to_xlsx_bytes(self):
        """
        导出为XLSX文件流
        
        Returns:
            BytesIO: 包含Excel文件的字节流
        """
        # 获取聚合数据
        aggregated_data = list(self._get_aggregated_data())
        
        if not aggregated_data:
            # 如果没有数据，创建空文件
            return self._create_empty_file()
        
        # 生成列标题和字段映射
        self.column_titles = self._generate_column_titles(aggregated_data)
        field_mapping = self._get_field_mapping(aggregated_data)
        
        # 创建工作簿和工作表
        self.workbook = openpyxl.Workbook()
        self.ws = self.workbook.active
        self.ws.title = "导出数据"
        
        # 设置标题行
        self._setup_worksheet()
        
        # 填充数据
        for row_num, item in enumerate(aggregated_data, 2):  # 从第2行开始
            row_data = [""] * len(self.column_titles)
            
            for field, value in item.items():
                if field in field_mapping:
                    col_index = field_mapping[field]
                    row_data[col_index] = self._convert_value(value)
            
            for col_num, value in enumerate(row_data, 1):
                self.ws.cell(row=row_num, column=col_num, value=value)
        
        # 自动调整列宽
        self._auto_adjust_column_width()
        
        # 转换为字节流
        file_stream = io.BytesIO()
        self.workbook.save(file_stream)
        file_stream.seek(0)
        
        return file_stream
    
    def _create_empty_file(self):
        """创建空文件"""
        self.workbook = openpyxl.Workbook()
        self.ws = self.workbook.active
        self.ws.title = "导出数据"
        self.ws['A1'] = "没有找到数据"
        
        file_stream = io.BytesIO()
        self.workbook.save(file_stream)
        file_stream.seek(0)
        
        return file_stream
    
    def get_download_response(self):
        """
        获取Django下载响应
        
        Returns:
            HttpResponse: 包含Excel文件的HTTP响应
        """
        from django.http import HttpResponse
        
        file_stream = self.export_to_xlsx_bytes()
        
        response = HttpResponse(
            file_stream.getvalue(),
            content_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
        )
        response['Content-Disposition'] = f'attachment; filename="{self.filename}"'
        
        return response


# 使用示例
def example_usage():
    """
    使用示例
    """
    from myapp.models import MyModel
    
    # 获取QuerySet
    queryset = MyModel.objects.filter(enabled=True).values()
    
    # 定义分组字段和求和字段
    group_by_fields = ['category', 'department']  # 按类别和部门分组
    sum_fields = ['amount', 'quantity']  # 对金额和数量求和
    
    # 创建导出器实例
    exporter = QuerySetToXlsxExporter(
        queryset=queryset,
        group_by_fields=group_by_fields,
        sum_fields=sum_fields,
        filename="销售数据汇总.xlsx"
    )
    
    # 获取下载响应
    return exporter.get_download_response()


# 高级用法：自定义导出
class CustomExporter(QuerySetToXlsxExporter):
    """自定义导出器，可以重写方法来定制行为"""
    
    def _convert_value(self, value):
        """自定义值转换逻辑"""
        # 调用父类方法进行基本转换
        converted = super()._convert_value(value)
        
        # 添加自定义逻辑
        if isinstance(value, (int, float)) and value == 0:
            return "零"
            
        return converted


# 在Django view中的使用示例
def export_data_view(request):
    """Django视图函数示例"""
    from myapp.models import SalesRecord
    
    # 获取查询参数
    start_date = request.GET.get('start_date')
    end_date = request.GET.get('end_date')
    
    # 构建QuerySet
    queryset = SalesRecord.objects.filter(enabled=True)
    
    if start_date and end_date:
        queryset = queryset.filter(create_date__range=[start_date, end_date])
    
    queryset = queryset.values()
    
    # 创建导出器
    exporter = QuerySetToXlsxExporter(
        queryset=queryset,
        group_by_fields=['product_category', 'sales_region'],
        sum_fields=['sales_amount', 'profit'],
        filename="销售报表.xlsx"
    )
    
    return exporter.get_download_response()